RSpec.shared_examples 'Msf::Exploit::SQLi::Common' do |sqli_class|
  let(:common_class) do
    sqli_class
  end
  before(:example) do
    # because vprint_status accesses framework, datastore and user_output
    allow_any_instance_of(common_class).to receive(:vprint_status).and_return(nil)
  end
  let(:datastore) { instance_double(::Msf::DataStore) }
  context 'Without opts' do
    let(:query_proc) do
      proc do |payload|
        payload[/'(.+?)'/, 1] || ''
      end
    end
    let(:sqli_obj) do
      common_class.new(datastore, {}, {}, &query_proc)
    end
    context('#run_sql') do
      queries = ["select concat(username,':',password) from users", "select 'hello'", 'select 1234 from users']
      query_results = [ ':', 'hello', '' ]
      queries.each_with_index do |query, i|
        it 'Should call vprint_status on run_sql' do
          expect(sqli_obj).to receive(:vprint_status).once
          expect(sqli_obj.run_sql(query)).to eql query_results[i]
        end
      end
    end

    context('#test_vulnerable') do
      it 'Should detect if the sqli object is expected to perform SQLi successfully' do
        expect(sqli_obj.test_vulnerable).to eql true
        allow(sqli_obj).to receive(:run_sql).and_return '<div id="articles"></div>'
        expect(sqli_obj.test_vulnerable).to eql false
      end
    end

    context('#dump_table_fields') do
      result_limit = rand(1..26)
      common_query = /^select.*password.*from.*maindb\.users\s*;?\s*?(?:#|--)?$/mi
      condition_query = /^select.*password.*from.*maindb\.users\s+where/mi
      limit_query = /^select.*password.*from.*maindb\.users\s+limit/mi

      # query without condition and limit
      it 'Should yield valid queries' do
        expect(sqli_obj).to receive(:run_sql).and_call_original
        expect(query_proc).to receive(:call).with(common_query).and_call_original
        sqli_obj.dump_table_fields('maindb.users', %w[password])
      end
      # query with condition string
      it 'Should yield valid queries when the user adds a condition' do
        expect(query_proc).to receive(:call).with(condition_query).and_call_original
        sqli_obj.dump_table_fields('maindb.users', %w[password], "username='admin'")
      end
      # query with limit
      it 'Should yield valid queries when the user adds a limit number' do
        expect(query_proc).to receive(:call).with(limit_query).and_call_original
        sqli_obj.dump_table_fields('maindb.users', %w[password], '', result_limit)
      end
    end
  end
  context 'truncation_length set' do
    let(:opts) do
      { truncation_length: rand(1..20) }
    end
    let(:query_proc) do
      proc do |payload|
        payload[/'(.+?)'/, 1] || ''
      end
    end
    let(:sqli_obj) do
      common_class.new(datastore, {}, {}, opts, &query_proc)
    end
    context '#truncated_query should act like run_sql' do
      let(:query_result) do
        'e951a99943ebe29c6fc425c7df2a0544,028cad8c0961163ef8401d3573b41d8e,b090e41c61a321c99bca94bdb26d1788'
      end
      let(:query_proc) do
        i = 0
        proc do |_payload|
          slice = query_result[i, opts[:truncation_length]]
          i += opts[:truncation_length]
          if slice.empty?
            i = 0
            ''
          else
            slice
          end
        end
      end
      let(:sqli_obj) do
        common_class.new({}, {}, {}, opts, &query_proc)
      end
      it 'Should concatenate the slices and return output like run_sql' do
        expect(sqli_obj.send(:truncated_query, 'select substr(username,^OFFSET^,' \
        "#{opts[:truncation_length]}) from users")).to eql query_result
      end
    end
    context 'call_function and dump_table_fields should call truncated_query instead of run_sql' do
      it '#dump_table_fields' do
        expect(sqli_obj).to receive(:truncated_query).and_call_original
        sqli_obj.dump_table_fields('users', %w[username])
      end
      it '#call_function' do
        # called by version(), current_user(), current_database() if sqli_obj responds to them
        expect(sqli_obj).to receive(:truncated_query)
        sqli_obj.send(:call_function, 'version()')
      end
    end
  end
  context 'custom encoder set' do
    let(:opts) do
      { encoder: { encode: 'reverse(^DATA^)', decode: :reverse.to_proc } }
    end
    let(:dump_data) do
      %w[
        ALL_PLUGINS APPLICABLE_ROLES CHARACTER_SETS CHECK_CONSTRAINTS COLLATIONS
        COLLATION_CHARACTER_SET_APPLICABILITY COLUMNS COLUMN_PRIVILEGES ENABLED_ROLES
      ]
    end
    let(:query_proc) do
      proc do
        # the server response should be encoded
        dump_data.map(&:reverse).join(',')
      end
    end
    let(:sqli_obj) do
      common_class.new(datastore, {}, {}, opts, &query_proc)
    end
    context '#initialize' do
      it 'should set the custom encoder' do
        expect(sqli_obj.instance_variable_get(:@encoder)).to eql opts[:encoder]
      end
    end

    context '#enum_table_names' do
      function_call = /reverse\(/i
      it 'Should apply the encoder correctly in the query, and use the decoder to retrieve decoded results' do
        expect(sqli_obj.instance_variable_get(:@query_proc)).to receive(:call).with(function_call).and_call_original
        expect(sqli_obj.enum_table_names).to eql dump_data
      end
    end
  end
end
